import logging

# These are the default prefixes from ComfyUI and can be found here:
# https://github.com/comfyanonymous/ComfyUI/blob/5352abc6d389570455776c457738db54367cd6cb/comfy/text_encoders/qwen_image.py#L17
# Qwen Image
DEFAULT_NO_IMAGES = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
# Qwen Image Edit
DEFAULT_WITH_IMAGES = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"


class CLIPSetQwenImagePrompt:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "clip": ("CLIP",),
                "template": (
                    "STRING",
                    {"multiline": True, "default": DEFAULT_NO_IMAGES},
                ),
            }
        }

    CATEGORY = "conditioning"
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "execute"

    def execute(self, clip, template):
        # Bail if not QwenImageTokenizer
        if clip.tokenizer.__class__.__name__ != "QwenImageTokenizer":
            return [clip]

        cloned = clip.clone()
        if hasattr(cloned.tokenizer, "llama_template"):
            cloned.tokenizer.llama_template = template
            logging.info("Patching Qwen Image encoder template.")

        return [cloned]


class CLIPSetQwenImageEditPrompt:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "clip": ("CLIP",),
                "template": (
                    "STRING",
                    {"multiline": True, "default": DEFAULT_WITH_IMAGES},
                ),
            }
        }

    CATEGORY = "conditioning"
    RETURN_TYPES = ("CLIP",)
    FUNCTION = "execute"

    def execute(self, clip, template):
        # Bail if not QwenImageTokenizer
        if clip.tokenizer.__class__.__name__ != "QwenImageTokenizer":
            return [clip]

        cloned = clip.clone()
        if hasattr(cloned.tokenizer, "llama_template_images"):
            cloned.tokenizer.llama_template_images = template
            logging.info("Patching Qwen Image Edit encoder template")

        return [cloned]


NODE_CLASS_MAPPINGS = {
    "CLIPSetQwenImagePrompt": CLIPSetQwenImagePrompt,
    "CLIPSetQwenImageEditPrompt": CLIPSetQwenImageEditPrompt,
}
